// Copyright (c) 2020 Presto Labs Pte. Ltd.
// Author: yuxuan

#pragma once

#include "coin2/strategy/util/moving_window.h"

namespace coin2::strategy::util {

enum OhlcType : int { OPEN = 0, HIGH = 1, LOW = 2, CLOSE = 3 };

template <typename T>
class MovingOhlc : public MovingWindow<T, T> {
 public:
  explicit MovingOhlc(int64_t time_window) : MovingWindow<T, T>(time_window) {}

  virtual const T& Get() = 0;

  void Update(int64_t ts) final {
    while (!MovingWindow<T, T>::Empty() &&
           ts - MovingWindow<T, T>::dq_.front().first > MovingWindow<T, T>::timewindow_) {
      MovingWindow<T, T>::dq_.pop_front();
    }
  }
};

template <typename T>
class MovingHigh : public MovingOhlc<T> {
 public:
  explicit MovingHigh(int64_t time_window) : MovingOhlc<T>(time_window) {}

  void Add(int64_t ts, const T& obj) final {
    MovingOhlc<T>::Update(ts);
    while (!MovingWindow<T, T>::Empty() && obj >= MovingWindow<T, T>::GetLastData()) {
      MovingWindow<T, T>::dq_.pop_back();
    }
    MovingWindow<T, T>::dq_.push_back({ts, obj});
  }

  const T& Get() final { return MovingWindow<T, T>::GetFirstData(); }
};

template <typename T>
class MemMovingHigh : public MovingOhlc<T> {
 public:
  explicit MemMovingHigh(int64_t time_window, int64_t merge_window)
      : MovingOhlc<T>(time_window), merge_window_(merge_window) {
    CHECK_GT(merge_window_, 0);
  }
  void Add(int64_t ts, const T& obj) final {
    int64_t interval = ts / merge_window_;
    ts = interval * merge_window_;
    MovingOhlc<T>::Update(ts);
    bool popped = false;
    while (!MovingWindow<T, T>::Empty() && obj >= MovingWindow<T, T>::GetLastData()) {
      MovingWindow<T, T>::dq_.pop_back();
      popped = true;
    }
    if (MovingWindow<T, T>::Empty() || popped ||
        MovingWindow<T, T>::GetLastTs() / merge_window_ < interval) {
      MovingWindow<T, T>::dq_.push_back({ts, obj});
    }
  }

  const T& Get() final { return MovingWindow<T, T>::GetFirstData(); }

 private:
  int64_t merge_window_;
};

template <typename T>
class MovingLow : public MovingOhlc<T> {
 public:
  explicit MovingLow(int64_t ts) : MovingOhlc<T>(ts) {}

  void Add(int64_t ts, const T& obj) final {
    MovingOhlc<T>::Update(ts);
    while (!MovingWindow<T, T>::Empty() && obj <= MovingWindow<T, T>::GetLastData()) {
      MovingWindow<T, T>::dq_.pop_back();
    }
    MovingWindow<T, T>::dq_.push_back({ts, obj});
  }

  const T& Get() final { return MovingWindow<T, T>::GetFirstData(); }
};

template <typename T>
class MemMovingLow : public MovingOhlc<T> {
 public:
  explicit MemMovingLow(int64_t ts, int64_t merge_window)
      : MovingOhlc<T>(ts), merge_window_(merge_window) {
    CHECK_GE(merge_window, 0);
  }

  void Add(int64_t ts, const T& obj) final {
    int64_t interval = ts / merge_window_;
    ts = interval * merge_window_;
    MovingOhlc<T>::Update(ts);
    bool popped = false;
    while (!MovingWindow<T, T>::Empty() && obj <= MovingWindow<T, T>::GetLastData()) {
      MovingWindow<T, T>::dq_.pop_back();
      popped = true;
    }
    if (MovingWindow<T, T>::Empty() || popped ||
        MovingWindow<T, T>::GetLastTs() / merge_window_ < interval) {
      MovingWindow<T, T>::dq_.push_back({ts, obj});
    }
  }

  const T& Get() final { return MovingWindow<T, T>::GetFirstData(); }

 private:
  int64_t merge_window_;
};

template <typename T>
class MovingOpen : public MovingOhlc<T> {
 public:
  explicit MovingOpen(int64_t ts) : MovingOhlc<T>(ts) {}

  void Add(int64_t ts, const T& obj) final {
    MovingOhlc<T>::Update(ts);
    MovingWindow<T, T>::dq_.push_back({ts, obj});
  }

  const T& Get() final { return MovingWindow<T, T>::GetFirstData(); }
};

template <typename T>
class MemMovingOpen : public MovingOhlc<T> {
 public:
  explicit MemMovingOpen(int64_t ts, int64_t merge_window)
      : MovingOhlc<T>(ts), merge_window_(merge_window) {
    CHECK_GT(merge_window_, 0);
  }

  void Add(int64_t ts, const T& obj) final {
    int64_t interval = ts / merge_window_;
    ts = interval * merge_window_;
    MovingOhlc<T>::Update(ts);
    CHECK_GE(interval, interval_id_);
    if (interval > interval_id_) {
      interval_id_ = interval;
      MovingWindow<T, T>::dq_.push_back({ts, obj});
    }
  }

  const T& Get() final { return MovingWindow<T, T>::GetFirstData(); }

 private:
  int64_t merge_window_;
  int64_t interval_id_{0};
};

template <typename T>
class MovingClose : public MovingOhlc<T> {
 public:
  explicit MovingClose(int64_t ts) : MovingOhlc<T>(ts) {}

  void Add(int64_t ts, const T& obj) final {
    MovingWindow<T, T>::dq_.clear();
    MovingWindow<T, T>::dq_.push_back({ts, obj});
  }

  const T& Get() final { return MovingWindow<T, T>::GetLastData(); }
};

template <typename T>
std::unique_ptr<MovingOhlc<T>>
DispatchMovingOhlc(const OhlcType ohlc_type, int64_t time_window, int64_t merge_time_window) {
  if (ohlc_type == OhlcType::OPEN) {
    if (time_window >= DEFAULT_MEM_EFFICIENT_THRESHOLD) {
      return std::move(std::make_unique<MemMovingOpen<T>>(time_window, merge_time_window));
    } else {
      return std::move(std::make_unique<MovingOpen<T>>(time_window));
    }
  } else if (ohlc_type == OhlcType::HIGH) {
    if (time_window >= DEFAULT_MEM_EFFICIENT_THRESHOLD) {
      return std::move(std::make_unique<MemMovingHigh<T>>(time_window, merge_time_window));
    } else {
      return std::move(std::make_unique<MovingHigh<T>>(time_window));
    }
  } else if (ohlc_type == OhlcType::LOW) {
    if (time_window >= DEFAULT_MEM_EFFICIENT_THRESHOLD) {
      return std::move(std::make_unique<MemMovingLow<T>>(time_window, merge_time_window));
    } else {
      return std::move(std::make_unique<MovingLow<T>>(time_window));
    }
  } else if (ohlc_type == OhlcType::CLOSE) {
    return std::move(std::make_unique<MovingClose<T>>(time_window));
  } else {
    return nullptr;
  }
}

}  // namespace coin2::strategy::util
